# -*- coding: utf-8 -*-

"""剑指 Offer II 066. 单词之和
实现一个 MapSum 类，支持两个方法，insert 和 sum：
MapSum() 初始化 MapSum 对象
void insert(String key, int val) 插入 key-val 键值对，字符串表示键 key ，整数表示值 val 。如果键 key 已经存在，那么原来的键值对将被替代成新的键值对。
int sum(string prefix) 返回所有以该前缀 prefix 开头的键 key 的值的总和。

示例：
输入：
inputs = ["MapSum", "insert", "sum", "insert", "sum"]
inputs = [[], ["apple", 3], ["ap"], ["app", 2], ["ap"]]
输出：
[null, null, 3, null, 5]

解释：
MapSum mapSum = new MapSum();
mapSum.insert("apple", 3);  
mapSum.sum("ap");           // return 3 (apple = 3)
mapSum.insert("app", 2);    
mapSum.sum("ap");           // return 5 (apple + app = 3 + 2 = 5)

提示：
1 <= key.length, prefix.length <= 50
key 和 prefix 仅由小写英文字母组成
1 <= val <= 1000
最多调用 50 次 insert 和 sum"""


class TrieNode():
    def __init__(self) -> None:
        self.children = [None]*26
        self.isEnd = False


class MapSum:
    """前缀树，搜索+递归遍历"""
    def __init__(self):
        """
        Initialize your data structure here.
        """
        self.trie = TrieNode()
        self.dikt = dict()

    def insert(self, key: str, val: int) -> None:
        if key not in self.dikt.keys():
            node = self.trie
            for ch in key:
                idx = ord(ch)-97
                if not node.children[idx]:
                    node.children[idx] = TrieNode()
                node = node.children[idx]
            node.isEnd = True
        self.dikt[key] = val

    def sum(self, prefix: str) -> int:
        node = self.trie
        for ch in prefix:
            idx = ord(ch)-97
            if node.children[idx]:
                node = node.children[idx]
            else:
                return 0
        
        total = 0
        def reverse(node, sword):
            nonlocal total
            if node.isEnd:
                total += self.dikt[sword]
            for ch, child in [(chr(idx+97), child) for idx, child in enumerate(node.children) if child]:
                reverse(child, sword+ch)
        
        reverse(node, prefix)
        return total


if __name__ == '__main__':
    mapSum = MapSum()
    # print(mapSum.insert("apple", 3))
    # print(mapSum.sum("ap"))
    # print(mapSum.insert("app", 2))
    # print(mapSum.sum("ap"))
    
    print(mapSum.insert("a", 3))
    print(mapSum.sum("ap"))
    print(mapSum.insert("b", 2))
    print(mapSum.sum("a"))
